package com.jjz.id.config;

import com.jjz.id.config.properties.JueAuthorizationProperties;

import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.JWSVerificationKeySelector;
import com.nimbusds.jose.proc.SecurityContext;
import java.security.KeyStore;
import com.nimbusds.jwt.proc.ConfigurableJWTProcessor;
import com.nimbusds.jwt.proc.DefaultJWTProcessor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.lang.Nullable;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.jwt.NimbusJwtDecoder;
import org.springframework.security.oauth2.jwt.NimbusJwtEncoder;
import org.springframework.security.oauth2.server.authorization.token.*;

/**
 * @author ZJL
 * @create 2025-05-28-9:39
 */
@Configuration(
        proxyBeanMethods = false
)
public class JwtConfiguration {
    @Bean
    public JWKSource<SecurityContext> jwkSource(JueAuthorizationProperties authorizationProperties) {
        try {
            ClassPathResource resource = new ClassPathResource(authorizationProperties.getJks().getPath());
            KeyStore jks = KeyStore.getInstance("jks");
            char[] pin = authorizationProperties.getJks().getPassword().toCharArray();
            jks.load(resource.getInputStream(), pin);
            RSAKey rsaKey = RSAKey.load(jks, authorizationProperties.getJks().getAlias(), pin);
            return (jwkSelector, securityContext) -> jwkSelector.select(new JWKSet(rsaKey));
        } catch (Exception e) {
            e.printStackTrace();
            System.out.println(e.getMessage());
        }
        return null;
    }
    @Bean
    OAuth2TokenGenerator<?> tokenGenerator(JWKSource<SecurityContext> jwkSource, @Nullable OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer, @Nullable OAuth2TokenCustomizer<OAuth2TokenClaimsContext> accCustomizer) {
        JwtGenerator jwtGenerator = new JwtGenerator(new NimbusJwtEncoder(jwkSource));
        if (jwtCustomizer != null) {
            jwtGenerator.setJwtCustomizer(jwtCustomizer);
        }

        OAuth2AccessTokenGenerator accessTokenGenerator = new OAuth2AccessTokenGenerator();
        if (accCustomizer != null) {
            accessTokenGenerator.setAccessTokenCustomizer(accCustomizer);
        }

        OAuth2RefreshTokenGenerator refreshTokenGenerator = new OAuth2RefreshTokenGenerator();
        return new DelegatingOAuth2TokenGenerator(new OAuth2TokenGenerator[]{jwtGenerator, accessTokenGenerator, refreshTokenGenerator});
    }

    @Bean
    JwtDecoder jwtDecoder(JWKSource<SecurityContext> jwkSource) {
        try {
            ConfigurableJWTProcessor<SecurityContext> jwtProcessor = new DefaultJWTProcessor();
            JWSVerificationKeySelector<SecurityContext> keySelector = new JWSVerificationKeySelector(JWSAlgorithm.RS256, jwkSource);
            jwtProcessor.setJWSKeySelector(keySelector);
            jwtProcessor.setJWTClaimsSetVerifier((claims, context) -> {
            });
            return new NimbusJwtDecoder(jwtProcessor);
        } catch (Exception e) {
            e.printStackTrace();
            System.out.println(e.getMessage());
        }
        return null;
    }
}
